#pragma once

#ifndef TORCH_HELPER
#define TORCH_HELPER
#include <torch/script.h>
#include <torch/torch.h>
#include <opencv2/opencv.hpp>

float generate_scale(const cv::Mat& image, const std::vector<int>& target_size);
float letterbox(const cv::Mat& input_image, cv::Mat& output_image, const std::vector<int>& target_size);
torch::Tensor xyxy2xywh(const torch::Tensor& x);
torch::Tensor xywh2xyxy(const torch::Tensor& x);
torch::Tensor nms(const torch::Tensor& bboxes, const torch::Tensor& scores, float iou_threshold);
torch::Tensor non_max_supperession(torch::Tensor& prediction, float conf_thres = 0.25, float iou_thres = 0.45f, int max_det = 300);
torch::Tensor clip_boxes(torch::Tensor& boxes, const std::vector<int>& shape);
torch::Tensor scale_boxes(const std::vector<int>& img1_shape, torch::Tensor& boxes, const std::vector<int>& img0_shape);
int run(cv::Mat& img, torch::jit::script::Module& yolo_model, std::vector<std::string>& classes, torch::Device& device);
#endif
